/***************************************************************************
** 版权所有:   Copyright (c) 2020-2030
** 文件名称:  showtask.c
** 当前版本:  v1.0
** 作    者:
** 完成日期: 2025-02-08
** 修改记录:
** 修改日期:
** 版本号  :
** 修改人  :
** 修改内容:
***************************************************************************/

/*****************************头文件****************************************/

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/sysfs.h>
#include <linux/kobject.h>
#include <linux/slab.h>
#include <linux/ctype.h>
#include <linux/uaccess.h>
#include <linux/pid.h>
#include <linux/fs.h>

#include "showtask.h"


/*****************************宏定义****************************************/
#define thread_group_leader(tsk) ((tsk)->group_leader)

/*****************************结构体或类型定义*******************************/

/*****************************全局变量****************************************/

/*****************************本地变量****************************************/

/*****************************函数或类声明****************************************/

/*****************************函数或类实现****************************************/

int showtasklist(const char *param)
{
    struct task_struct *task;
    int iTaskCount = 0;

    for_each_process(task)
    {
        printk(KERN_INFO "Process: %s (PID: %d)\n", task->comm, task->pid);
        iTaskCount++;
    }
    printk(KERN_INFO "task count:%d\n", iTaskCount);

    return 0;
}


int print_task_info(const char *param)
{
    struct task_struct *task;
    int iTaskCount = 0;

    for_each_process(task)
    {
        pid_t pid = task->pid;
        pid_t tgid = task->tgid;
        int prio = task->prio;
        int static_prio = task->static_prio;
        int normal_prio = task->normal_prio;

        printk(KERN_INFO "Process: %s (PID: %d, TGID: %d, Prio: %d, Static Prio: %d, Normal Prio: %d)\n",
               task->comm, pid, tgid, prio, static_prio, normal_prio);
        iTaskCount++;
    }
    printk(KERN_INFO "task count:%d\n", iTaskCount);
    return 0;
}


// 设置所有进程的优先级
static int set_all_task_nice(const char *param)
{
    int nice_value = 0;
    if ((sscanf(param, "%d", &nice_value) >= 1) && (nice_value >= -20 && nice_value <= 19))
    {
        struct task_struct *task;
        for_each_process(task)
        {
            if (task->pid != 0)
            { // 跳过 init 进程
                set_user_nice(task, nice_value);
                printk(KERN_INFO "Set nice value of process %s (PID: %d) to %d\n", task->comm, task->pid, nice_value);
            }
        }
    }
    else
    {
        printk(KERN_ERR "Invalid nice value\n");
    }

    return 0;
}

// 设置指定用户的进程的优先级
static int set_user_nice_by_uid(const char *param)
{
    int uid, nice_value;
    bool flag = false;

    if (sscanf(param, "%d %d", &uid, &nice_value) == 2)
    {
        struct task_struct *task;
        kuid_t kuid = make_kuid(current_user_ns(), uid);

        if (!uid_valid(kuid))
        {
            printk(KERN_ERR "Invalid UID: %d\n", uid);
            return -1;
        }

        if (nice_value < -20 || nice_value > 19)
        {
            printk(KERN_ERR "Invalid nice value: %d\n", nice_value);
            return -1;
        }

        for_each_process(task)
        {
            if (task->pid != 0 && uid_eq(task_uid(task), kuid))
            { // 跳过 init 进程
                set_user_nice(task, nice_value);
                printk(KERN_INFO "Set nice value of process %s (PID: %d, UID: %d) to %d\n", task->comm, task->pid, uid, nice_value);
                flag = true;
            }
        }
        if (!flag)
        {
            printk(KERN_ERR "No process found with UID: %d\n", uid);
        }
    }
    else
    {
        printk(KERN_ERR "Invalid input format. Expected: setniceuser <uid> <nice_value>\n");
    }

    return 0;
}
// 将调度策略转换为字符串
static const char *get_policy_str(int policy)
{
    switch (policy)
    {
    case SCHED_NORMAL:
        return "SCHED_NORMAL";
    case SCHED_FIFO:
        return "SCHED_FIFO";
    case SCHED_RR:
        return "SCHED_RR";
    case SCHED_BATCH:
        return "SCHED_BATCH";
    case SCHED_IDLE:
        return "SCHED_IDLE";
    case SCHED_DEADLINE:
        return "SCHED_DEADLINE";
    default:
        return "UNKNOWN";
    }
} 


int showallthread(const char *param)
{
    struct task_struct *task, *thread;
    int thread_count = 0;
    char buffer[1024]; // 用于存储打印信息的缓冲区

    // 打印标题
    snprintf(buffer, sizeof(buffer), "%-10s%-10s%-20s%-20s%-10s%-6s%s\n", 
             "PID", "TID", "Process Name", "Thread Name", "Priority", "Nice", "Policy");
    printk(KERN_INFO "%s", buffer);

    rcu_read_lock();
    for_each_process(task)
    {
        for_each_thread(task, thread)
        {
            int prio = thread->prio;
            int nice = task_nice(thread);
            int policy = thread->policy;
            const char *policy_str = get_policy_str(policy);

            // 将线程信息存储到缓冲区
            snprintf(buffer, sizeof(buffer), "%-10d%-10d%-20s%-20s%-10d%-6d%s\n",
                     task_pid_nr(task), task_pid_nr(thread), 
                     task->comm, thread->comm, prio, nice, policy_str);
            printk(KERN_INFO "%s", buffer);

            thread_count++;
        }
    }
    rcu_read_unlock();

    // 打印线程总数
    snprintf(buffer, sizeof(buffer), "Total number of threads: %d\n", thread_count);
    printk(KERN_INFO "%s", buffer);

    return 0;
}


#if 1
// 设置线程优先级的函数
int set_thread_priority(const char *param)
{
    struct task_struct *thread;
    struct pid *pid_struct;
    int tid, new_priority, old_priority;

    // 检查输入是否为空
    if (!param || strlen(param) == 0)
    {
        printk(KERN_INFO "Invalid input format. Expected: setthreadpriority <tid> <priority>\n");
        return -EINVAL;
    }

    printk(KERN_INFO "Setting thread priority: %s\n", param);

    // 使用 sscanf 解析参数
    if (sscanf(param, "%d %d", &tid, &new_priority) != 2)
    {
        printk(KERN_INFO "Invalid input format. Expected: setthreadpriority <tid> <priority>\n");
        return -EINVAL;
    }

    // 查找指定的线程
    pid_struct = find_get_pid(tid);
    if (!pid_struct)
    {
        printk(KERN_INFO "Thread with TID %d not found\n", tid);
        return -ESRCH;
    }

    thread = pid_task(pid_struct, PIDTYPE_PID);
    if (!thread)
    {
        printk(KERN_INFO "Thread with TID %d not found\n", tid);
        put_pid(pid_struct);
        return -ESRCH;
    }

    put_pid(pid_struct);

    // 检查优先级范围
    if (new_priority < MIN_NICE || new_priority > MAX_NICE)
    {
        printk(KERN_INFO "Invalid priority value: %d,in %d~%d\n", new_priority, MIN_NICE, MAX_NICE);
        return -EINVAL;
    }

    // 获取当前的 nice 值
    old_priority = task_nice(thread);

    // 设置新的优先级
    set_user_nice(thread, new_priority);

    // 获取进程名称（线程组领导者名称）和线程名称
    const char *process_name = thread_group_leader(thread)->comm;
    const char *thread_name = thread->comm;

    // 打印进程名称、线程名称以及修改前后的优先级
    printk(KERN_INFO "Thread %d (Process: %s, Thread: %s) priority changed from %d to %d\n",
           tid, process_name, thread_name, old_priority, new_priority);

    return 0;
}


#else

#include <linux/sched.h>
#include <linux/sched/signal.h>
#include <linux/sched/prio.h>
#include <linux/pid.h>
#include <linux/unistd.h>
#include <linux/kernel.h>
#include <linux/slab.h>
#include <linux/string.h>
#include <linux/module.h>
#include <linux/init.h>
#include <linux/sched/signal.h>
#include <linux/sched/rt.h>
#include <linux/sched/deadline.h>
// 设置线程优先级的函数
int set_thread_priority(const char *param)
{
    struct task_struct *thread;
    struct pid *pid_struct;
    int tid, new_priority, old_priority;
    struct sched_param param_struct;
    int current_policy;

    // 检查输入是否为空
    if (!param || strlen(param) == 0)
    {
        printk(KERN_INFO "Invalid input format. Expected: setthreadpriority <tid> <priority>\n");
        return -EINVAL;
    }

    printk(KERN_INFO "Setting thread priority: %s\n", param);

    // 使用 sscanf 解析参数
    if (sscanf(param, "%d %d", &tid, &new_priority) != 2)
    {
        printk(KERN_INFO "Invalid input format. Expected: setthreadpriority <tid> <priority>\n");
        return -EINVAL;
    }

    // 查找指定的线程
    pid_struct = find_get_pid(tid);
    if (!pid_struct)
    {
        printk(KERN_INFO "Thread with TID %d not found\n", tid);
        return -ESRCH;
    }

    thread = pid_task(pid_struct, PIDTYPE_PID);
    if (!thread)
    {
        printk(KERN_INFO "Thread with TID %d not found\n", tid);
        put_pid(pid_struct);
        return -ESRCH;
    }

    put_pid(pid_struct);

    // 检查优先级范围
    if (new_priority < MIN_RT_PRIO || new_priority > MAX_RT_PRIO)
    {
        printk(KERN_INFO "Invalid priority value: %d, in %d~%d\n", new_priority, MIN_RT_PRIO, MAX_RT_PRIO);
        return -EINVAL;
    }
    // 获取当前的 nice 值
    old_priority = task_nice(thread);

    // 获取当前的调度策略
    current_policy = thread->policy;

    // 设置新的优先级
    param_struct.sched_priority = new_priority;
    sched_setscheduler_nocheck(thread, current_policy, &param_struct);

    // 获取进程名称（线程组领导者名称）和线程名称
    const char *process_name = thread_group_leader(thread)->comm;
    const char *thread_name = thread->comm;

    // 打印进程名称、线程名称以及修改前后的优先级
    printk(KERN_INFO "Thread %d (Process: %s, Thread: %s) priority changed from %d to %d\n",
           tid, process_name, thread_name, old_priority, new_priority);

    return 0;
}

#endif
